#ifndef KALMAN_FILTER_BASE_HPP
#define KALMAN_FILTER_BASE_HPP

#include <Eigen/Dense>
#include <vector>
#include <stdexcept> // For std::runtime_error
#include <iostream> // For debugging

// Forward declaration if StateType and MeasurementType are complex
// For simplicity, we'll assume Eigen::VectorXf for now, but it could be templated further.

class KalmanFilterBase {
public:
    // state_dim: 状态向量的维度 (多少个需要估计的变量)
    // measurement_dim: 测量向量的维度 (从传感器直接读取多少个值)
    KalmanFilterBase(int state_dim, int measurement_dim) :
        state_dim_(state_dim),
        measurement_dim_(measurement_dim),
        is_initialized_(false) {
        if (state_dim <= 0 || measurement_dim <= 0) {
            throw std::invalid_argument("State and measurement dimensions must be positive.");
        }

        x_ = Eigen::VectorXf::Zero(state_dim_);
        P_ = Eigen::MatrixXf::Identity(state_dim_, state_dim_);
        F_ = Eigen::MatrixXf::Identity(state_dim_, state_dim_); // Default: state doesn't change without input
        H_ = Eigen::MatrixXf::Zero(measurement_dim_, state_dim_);
        Q_ = Eigen::MatrixXf::Identity(state_dim_, state_dim_) * 0.01f; // Default small process noise
        R_ = Eigen::MatrixXf::Identity(measurement_dim_, measurement_dim_) * 0.1f; // Default measurement noise
        I_ = Eigen::MatrixXf::Identity(state_dim_, state_dim_);

        // Default H: assume direct measurement of the first measurement_dim states
        // This often needs to be customized by the user or a derived class.
        for (int i = 0; i < std::min(state_dim_, measurement_dim_); ++i) {
            H_(i, i) = 1.0f;
        }

        // Initialize max_state_changes_ if needed (e.g., with a very large value or disabled)
        max_state_changes_ = Eigen::VectorXf::Constant(state_dim_, std::numeric_limits<float>::max());
    }

    virtual ~KalmanFilterBase() = default;

    // Initialize the filter state
    virtual void init(const Eigen::VectorXf& initial_state, const Eigen::MatrixXf& initial_covariance) {
        if (initial_state.size() != state_dim_ ||
            initial_covariance.rows() != state_dim_ ||
            initial_covariance.cols() != state_dim_) {
            throw std::invalid_argument("Initial state or covariance dimensions mismatch.");
        }
        x_ = initial_state;
        P_ = initial_covariance;
        is_initialized_ = true;
    }

    // Initialize with just initial state, default covariance (Identity * small_value)
    virtual void init(const Eigen::VectorXf& initial_state) {
        if (initial_state.size() != state_dim_) {
            throw std::invalid_argument("Initial state dimensions mismatch.");
        }
        x_ = initial_state;
        P_ = Eigen::MatrixXf::Identity(state_dim_, state_dim_) * 1.0f; // Or some other default
        is_initialized_ = true;
    }


    // Prediction step
    virtual void predict() {
        if (!is_initialized_) {
            // std::cerr << "Warning: Kalman filter not initialized before predict." << std::endl;
            // Or throw: throw std::runtime_error("Kalman filter not initialized.");
            return; // Or initialize with a default state if appropriate
        }
        x_ = F_ * x_;
        P_ = F_ * P_ * F_.transpose() + Q_;
    }

    // Update step with new measurement
    virtual void update(const Eigen::VectorXf& measurement) {
        if (!is_initialized_) {
            // If not initialized, use the first measurement to initialize
            // This is a common practice for simple KFs
            std::cout << "Kalman filter not initialized. Initializing with first measurement." << std::endl;
            Eigen::VectorXf initial_state = Eigen::VectorXf::Zero(state_dim_);
            // Try to map measurement to state if H allows (e.g., H is identity for some parts)
            // This is a heuristic and might need to be more sophisticated
            if (H_.cols() == state_dim_ && H_.rows() == measurement_dim_) {
                 // A simple pseudo-inverse approach if H is not square or not full rank for direct inversion
                // This is a placeholder; a proper pseudo-inverse or specific logic is better.
                // For simplicity, if state_dim == measurement_dim and H is identity, this works.
                if (state_dim_ == measurement_dim_ && H_.isApprox(Eigen::MatrixXf::Identity(state_dim_, state_dim_))) {
                    initial_state = measurement;
                } else {
                    // More complex mapping or use a default zero state
                    // For example, if H selects first N components:
                    for(int i=0; i<std::min(state_dim_, measurement_dim_); ++i) {
                        if(H_(i,i) == 1.0f) initial_state(i) = measurement(i);
                    }
                }
            }
            init(initial_state); // Initialize with (mapped) measurement and default P
        }

        Eigen::VectorXf y = measurement - H_ * x_;          // Measurement residual (innovation)
        Eigen::MatrixXf S = H_ * P_ * H_.transpose() + R_;  // Residual covariance
        Eigen::MatrixXf K = P_ * H_.transpose() * S.inverse(); // Kalman gain

        Eigen::VectorXf x_new = x_ + K * y;

        // Apply max state change constraint if enabled
        for (int i = 0; i < state_dim_; ++i) {
            float change = x_new(i) - x_(i);
            if (std::abs(change) > max_state_changes_(i)) {
                x_new(i) = x_(i) + (change > 0 ? max_state_changes_(i) : -max_state_changes_(i));
            }
        }
        x_ = x_new;
        P_ = (I_ - K * H_) * P_;
    }

    // Get current state estimate
    Eigen::VectorXf getState() const {
        return x_;
    }

    // Get current state covariance
    Eigen::MatrixXf getCovariance() const {
        return P_;
    }

    bool isInitialized() const {
        return is_initialized_;
    }

    // --- Setters for matrices and parameters ---

    // Set the entire state transition matrix F
    void setStateTransitionMatrix(const Eigen::MatrixXf& F) {
        if (F.rows() != state_dim_ || F.cols() != state_dim_) {
            throw std::invalid_argument("F matrix dimensions mismatch.");
        }
        F_ = F;
    }

    // Set the entire measurement matrix H
    void setMeasurementMatrix(const Eigen::MatrixXf& H) {
        if (H.rows() != measurement_dim_ || H.cols() != state_dim_) {
            throw std::invalid_argument("H matrix dimensions mismatch.");
        }
        H_ = H;
    }

    // Set process noise covariance Q (as a full matrix)
    void setProcessNoiseCovariance(const Eigen::MatrixXf& Q) {
        if (Q.rows() != state_dim_ || Q.cols() != state_dim_) {
            throw std::invalid_argument("Q matrix dimensions mismatch.");
        }
        Q_ = Q;
    }

    // Set process noise for each state variable (diagonal Q)
    void setProcessNoise(const Eigen::VectorXf& q_diag_noises) {
        if (q_diag_noises.size() != state_dim_) {
            throw std::invalid_argument("Process noise vector dimension mismatch.");
        }
        Q_ = Eigen::MatrixXf::Zero(state_dim_, state_dim_);
        for (int i = 0; i < state_dim_; ++i) {
            Q_(i, i) = q_diag_noises(i);
        }
    }

    // Set a single process noise value for all state variables (diagonal Q)
    void setProcessNoise(float q_noise) {
        if (q_noise < 0) throw std::invalid_argument("Noise value cannot be negative.");
        Q_ = Eigen::MatrixXf::Identity(state_dim_, state_dim_) * q_noise;
    }


    // Set measurement noise covariance R (as a full matrix)
    void setMeasurementNoiseCovariance(const Eigen::MatrixXf& R) {
        if (R.rows() != measurement_dim_ || R.cols() != measurement_dim_) {
            throw std::invalid_argument("R matrix dimensions mismatch.");
        }
        R_ = R;
    }

    // Set measurement noise for each measurement channel (diagonal R)
    void setMeasurementNoise(const Eigen::VectorXf& r_diag_noises) {
        if (r_diag_noises.size() != measurement_dim_) {
            throw std::invalid_argument("Measurement noise vector dimension mismatch.");
        }
        R_ = Eigen::MatrixXf::Zero(measurement_dim_, measurement_dim_);
        for (int i = 0; i < measurement_dim_; ++i) {
            R_(i, i) = r_diag_noises(i);
        }
    }

    // Set a single measurement noise value for all measurement channels (diagonal R)
    void setMeasurementNoise(float r_noise) {
        if (r_noise < 0) throw std::invalid_argument("Noise value cannot be negative.");
        R_ = Eigen::MatrixXf::Identity(measurement_dim_, measurement_dim_) * r_noise;
    }


    // Set maximum state change limits for each state variable
    void setMaxStateChanges(const Eigen::VectorXf& max_changes) {
        if (max_changes.size() != state_dim_) {
            throw std::invalid_argument("Max state changes vector dimension mismatch.");
        }
        for(int i=0; i<state_dim_; ++i) {
            if(max_changes(i) < 0) throw std::invalid_argument("Max change value cannot be negative.");
        }
        max_state_changes_ = max_changes;
    }

    // Set a single maximum state change value for all state variables
    void setMaxStateChange(float max_change) {
        if (max_change < 0) throw std::invalid_argument("Max change value cannot be negative.");
        max_state_changes_ = Eigen::VectorXf::Constant(state_dim_, max_change);
    }


protected:
    int state_dim_;
    int measurement_dim_;
    bool is_initialized_;

    Eigen::VectorXf x_; // State estimate: [x, y, vx, vy, ...]
    Eigen::MatrixXf P_; // State covariance matrix
    Eigen::MatrixXf F_; // State transition matrix
    Eigen::MatrixXf H_; // Measurement matrix
    Eigen::MatrixXf Q_; // Process noise covariance
    Eigen::MatrixXf R_; // Measurement noise covariance
    Eigen::MatrixXf I_; // Identity matrix

    Eigen::VectorXf max_state_changes_; // Max allowed change per update for each state component
};

#endif // KALMAN_FILTER_BASE_HPP